import rclpy
from rclpy.node import Node
from nav_msgs.msg import Path
from geometry_msgs.msg import Twist, PoseStamped
from tf2_ros import TransformException
from tf2_ros.buffer import Buffer
from tf2_ros.transform_listener import TransformListener

class ROSInterface:
    def __init__(self, node, callback_group):
        self.node = node
        self.callback_group = callback_group
        
        # TF2设置
        self.tf_buffer = Buffer()
        self.tf_listener = TransformListener(self.tf_buffer, node)
        
        # 创建订阅者
        self.path_sub = node.create_subscription(
            Path,
            'path',
            self._path_callback,
            10,
            callback_group=callback_group
        )
        
        # 创建发布者
        self.cmd_vel_pub = node.create_publisher(
            Twist,
            'cmd_vel',
            10
        )
        
        # 回调函数
        self.path_callback = None
        
        # 获取参数
        self.base_frame = node.get_parameter('base_frame').value
        self.global_frame = node.get_parameter('global_frame').value
    
    def set_path_callback(self, callback):
        """设置路径消息回调函数"""
        self.path_callback = callback
    
    def _path_callback(self, path_msg):
        """路径消息内部回调函数"""
        if self.path_callback:
            self.path_callback(path_msg)
    
    def get_current_pose(self):
        """获取机器人当前位姿"""
        try:
            # 获取从全局坐标系到机器人基坐标系的变换
            transform = self.tf_buffer.lookup_transform(
                self.global_frame,
                self.base_frame,
                rclpy.time.Time()
            )
            
            # 提取位置
            x = transform.transform.translation.x
            y = transform.transform.translation.y
            
            # 提取朝向（四元数转欧拉角）
            q = transform.transform.rotation
            siny_cosp = 2 * (q.w * q.z + q.x * q.y)
            cosy_cosp = 1 - 2 * (q.y * q.y + q.z * q.z)
            theta = np.arctan2(siny_cosp, cosy_cosp)
            
            return x, y, theta
        except TransformException as ex:
            self.node.get_logger().error(f'无法获取变换: {ex}')
            return None, None, None
    
    def publish_velocity(self, linear_vel, angular_vel):
        """发布速度命令"""
        twist = Twist()
        twist.linear.x = linear_vel
        twist.angular.z = angular_vel
        self.cmd_vel_pub.publish(twist)